#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import pymongo
import re
import click
from tqdm import tqdm

DEFAULT_MONGO_URI = "mongodb://mongo-wepd.wepd:27017/wepd"
DEFAULT_VARIANT_COLLECTION = "variant"


@click.command()
@click.option("--mongo_uri", "-m", help="MongoDB URI", required=True, type=str, default=DEFAULT_MONGO_URI,
              show_default=True)
@click.option("--variant", "-v", help="Variant collection", required=True, type=str, default=DEFAULT_VARIANT_COLLECTION,
              show_default=True)
@click.argument("vcf_path", type=click.Path(exists=True))
def wepd_annovar2mongo(mongo_uri: str, variant: str, annovar_path: str):
    """
    import WEPD annovar annotation file into MongoDB
    :param mongo_uri: MongoDB connection URI
    :param variant: MongoDB variant collection
    :param annovar_path: input annovar file path
    :return: none
    """
    client = pymongo.MongoClient(mongo_uri)
    db = client.get_default_database()
    collection = db[variant]
    current_chr = ""

    infile = open(annovar_path, "r")
    for line in tqdm(infile, desc="line"):
        line = line.rstrip("\n")
        splitline = line.split("\t")
        if splitline[0] != "intergenic":
            gene = splitline[1]
            gene = re.sub("\(.*\)", "", gene)
            genes = re.split("[,|;]", gene)
            chrom = splitline[2].strip()
            pos = int(splitline[3].strip())

            if current_chr != chrom:
                current_chr = chrom
                tqdm.write(current_chr)

            collection.update_many({"chrom_id": chrom, "position": pos},
                                   {"$set": {"gene": genes, "gene_region": splitline[0]}}, upsert=False)

    infile.close()
    client.close()


if __name__ == '__main__':
    # wepd_annovar2mongo("/Users/genesis/Seafile/Work/华表计划/HuaBiao.SNP.only.recode.variant_function")
    wepd_annovar2mongo()
